// Copyright 2025 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef MEDIAPIPE_FRAMEWORK_API3_CALCULATOR_CONTEXT_H_
#define MEDIAPIPE_FRAMEWORK_API3_CALCULATOR_CONTEXT_H_

#include <memory>
#include <string>
#include <type_traits>

#include "absl/log/absl_check.h"
#include "absl/log/absl_log.h"
#include "mediapipe/framework/api3/contract.h"
#include "mediapipe/framework/api3/internal/contract_fields.h"
#include "mediapipe/framework/api3/internal/contract_to_tuple.h"
#include "mediapipe/framework/api3/internal/port_base.h"
#include "mediapipe/framework/api3/internal/specializers.h"
#include "mediapipe/framework/api3/one_of.h"
#include "mediapipe/framework/api3/packet.h"
#include "mediapipe/framework/calculator_context.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/graph_service.h"
#include "mediapipe/framework/resources.h"
#include "mediapipe/framework/timestamp.h"

namespace mediapipe::api3 {

// Calculator context specialized for a specific node.
//
// (Side)Inputs, (Side)Outputs and Options must be accessed through this
// context. (E.g. cc.input, cc.options etc.)
template <typename NodeT>
class CalculatorContext : public NodeT::template Contract<ContextSpecializer> {
 public:
  explicit CalculatorContext(mediapipe::CalculatorContext& generic_context) {
    holder_ = std::make_unique<internal_port::CalculatorContextHolder>();
    holder_->context = &generic_context;
    typename NodeT::template Contract<ContextSpecializer>* ptr = this;
    auto field_ptrs = ContractToFieldPtrTuple(*ptr);
    std::apply(
        [&](auto&... args) {
          ((internal_port::SetCalculatorContextHolder(*args, *holder_)), ...);
        },
        field_ptrs);
  }

  // Returns the current input timestamp.
  Timestamp InputTimestamp() const {
    return holder_->context->InputTimestamp();
  }

  // Returns a requested service binding.
  //
  // NOTE: you can request service in `UpdateContract` when defining a node or
  //   when implemenating it as a calculator.
  template <typename T>
  ServiceBinding<T> Service(const GraphService<T>& service) {
    return holder_->context->Service(service);
  }

  // Gets interface to access resources (file system, assets, etc.) from
  // calculators.
  //
  // NOTE: this is the preferred way to access resources from subgraphs and
  // calculators as it allows for fine grained per graph configuration.
  //
  // Resources can be configured by setting a custom `kResourcesService` graph
  // service on `CalculatorGraph`. The default resources service can be created
  // and reused through `CreateDefaultResources`.
  const Resources& GetResources() const {
    return holder_->context->GetResources();
  }

  // Returns calculator node name generated by the framework.
  // NOTE: the framework generates this depending on various sources
  //   (e.g. node.name/calculator). Should mostly used for debugging/profiling.
  const std::string& NodeName() const { return holder_->context->NodeName(); }

  mediapipe::CalculatorContext& GetGenericContext() {
    return *holder_->context;
  }

 private:
  void Reset(mediapipe::CalculatorContext& generic_context) {
    if (holder_->context != nullptr) {
      ABSL_LOG(DFATAL) << "Object must be cleared before resetting.";
    }
    holder_->context = &generic_context;
  }

  void Clear() {
    if (holder_->context == nullptr) {
      ABSL_LOG(DFATAL) << "Object has been already cleared.";
    }
    holder_->context = nullptr;
  }

  std::unique_ptr<internal_port::CalculatorContextHolder> holder_;

  // To Reset/Clear the context.
  template <typename N, typename I>
  friend class Calculator;
};

// +----------------------------------------------------------------------+
// |                                                                      |
// |   Specializations of (Side)Input/Output and Options for CONTEXT.     |
// |                                                                      |
// +----------------------------------------------------------------------+

template <typename PayloadT>
class Input<ContextSpecializer, PayloadT>
    : public internal_port::Port<ContextSpecializer, InputStreamField> {
 public:
  using Payload = PayloadT;
  using internal_port::Port<ContextSpecializer, InputStreamField>::Port;

  // Checks if payload can be accessed - valid input & non empty packet.
  explicit operator bool() const {
    auto id = holder_->context->Inputs().GetId(Tag(), Index());
    return id.IsValid() &&
           !holder_->context->Inputs().Get(id).Value().IsEmpty();
  }

  // Returns the payload of the packet for this particular input.
  //
  // NOTE: Dies if input packet is missing, input must be checked before
  //   accessing the payload, e.g. `RET_CHECK(cc.input)`
  const PayloadT& GetOrDie() const {
    return holder_->context->Inputs()
        .Get(Tag(), Index())
        .Value()
        .template Get<PayloadT>();
  }

  // Returns the packet for this particular input or empty packet if it's
  // missing.
  mediapipe::api3::Packet<PayloadT> Packet() const {
    auto maybe_packet = WrapLegacyPacket<PayloadT>(
        holder_->context->Inputs().Get(Tag(), Index()).Value());
    // Framework guarantees correct type (PayloadT) of the Input.
    ABSL_CHECK_OK(maybe_packet);
    return std::move(maybe_packet).value();
  }
};

template <typename... PayloadTs>
class Input<ContextSpecializer, OneOf<PayloadTs...>>
    : public internal_port::Port<ContextSpecializer, InputStreamField> {
 public:
  using Payload = OneOf<PayloadTs...>;
  using internal_port::Port<ContextSpecializer, InputStreamField>::Port;

  // Checks if payload can be accessed - valid input & non empty packet.
  explicit operator bool() const {
    auto id = holder_->context->Inputs().GetId(Tag(), Index());
    return id.IsValid() &&
           !holder_->context->Inputs().Get(id).Value().IsEmpty();
  }

  // Checks if input holds a packet of a specific type. Returns false if the
  // payload can't be accessed (i.e. it is not required to check the input
  // first).
  template <typename T>
  bool Has() const {
    static_assert((std::is_same_v<T, PayloadTs> || ...),
                  "Stream type must be one of the types in OneOf");
    return holder_->context->Inputs()
        .Get(Tag(), Index())
        .Value()
        .template ValidateAsType<T>()
        .ok();
  }

  // Returns the payload of the requested type for this particular input.
  //
  // NOTE: Dies if input packet is missing or of the wrong type. Input must be
  //   checked before accessing the payload, e.g. `RET_CHECK(cc.input.Has<T>())`
  template <typename T>
  const T& GetOrDie() const {
    static_assert((std::is_same_v<T, PayloadTs> || ...),
                  "Requested type must be one of the types in OneOf");
    return holder_->context->Inputs()
        .Get(Tag(), Index())
        .Value()
        .template Get<T>();
  }

  // Returns the packet of the requested type for this particular input.
  //
  // NOTE: Dies if input packet is missing or of the wrong type. Input must be
  //   checked before accessing the packet, e.g. `cc.input.Has<T>()`.
  template <typename T>
  mediapipe::api3::Packet<T> PacketOrDie() const {
    static_assert((std::is_same_v<T, PayloadTs> || ...),
                  "Requested packet type must be one of the types in OneOf");
    auto maybe_packet = WrapLegacyPacket<T>(
        holder_->context->Inputs().Get(Tag(), Index()).Value());
    // User is required to supply the correct T.
    ABSL_CHECK_OK(maybe_packet);
    return std::move(maybe_packet).value();
  }

  // Visits the value in the input with the given visitor lambdas.
  // There must be one lambda provided for each type in PayloadTs.
  //
  // NOTE: Dies if input packet is missing, input must be checked before
  //   accessing the payload, e.g. `RET_CHECK(cc.input)`
  //
  // Example:
  // ```
  //   input.VisitOrDie(
  //     [&](const TypeA& a) { ... },
  //     [&](const TypeB& b) { ... }
  //   );
  // ```
  template <int&... DoNotSpecify, class... F,
            std::enable_if_t<(sizeof...(F) > 1), bool> = true>
  auto VisitOrDie(F&&... visitors) const;

  // Visits the value in the input with the given visitor lambda.
  // There must be a single visitor provided that must be callable with all
  // types in PayloadTs.
  //
  // NOTE: Dies if input packet is missing, input must be checked before
  //   accessing the payload, e.g. `RET_CHECK(cc.input)`
  //
  // Example:
  // ```
  //   input.VisitOrDie(
  //     absl::Overload(
  //       [&](const TypeA& a) { ... },
  //       // Fallback.
  //       [&](const auto& b) { ... })
  //   );
  // ```
  template <int&... DoNotSpecify, class Visitor>
  auto VisitOrDie(Visitor&& visitor) const;

  // Visits the value in the input with the given visitor lambda, passing it
  // as `Packet<T>` rather than `const T&`.
  // There must be a single visitor provided that must be callable with all
  // `Packet<PayloadTs>` types.
  //
  // NOTE: Dies if input packet is missing, input must be checked before
  //   accessing the payload, e.g. `RET_CHECK(cc.input)`
  //
  // Example:
  // ```
  //   input.VisitAsPacketOrDie(
  //     absl::Overload(
  //       [&](Packet<TypeA> a) { ... },
  //       // Fallback.
  //       [&]<typename T>(Packet<T> b) { ... })
  //   );
  // ```
  template <int&... DoNotSpecify, class Visitor>
  auto VisitAsPacketOrDie(Visitor&& visitor) const;
};

template <typename PayloadT>
class SideInput<ContextSpecializer, PayloadT>
    : public internal_port::Port<ContextSpecializer, InputSidePacketField> {
 public:
  using Payload = PayloadT;
  using internal_port::Port<ContextSpecializer, InputSidePacketField>::Port;

  // Checks if payload can be accessed - valid side input & non empty packet.
  virtual explicit operator bool() const {
    auto id = holder_->context->InputSidePackets().GetId(Tag(), Index());
    return id.IsValid() &&
           !holder_->context->InputSidePackets().Get(id).IsEmpty();
  }

  // Returns the payload of the packet for this particular input.
  //
  // NOTE: Dies if input packet is missing, side input must be checked before
  //   accessing the payload, e.g. `RET_CHECK(cc.input)`
  const PayloadT& GetOrDie() const {
    return holder_->context->InputSidePackets()
        .Get(Tag(), Index())
        .template Get<PayloadT>();
  }

  mediapipe::api3::Packet<PayloadT> Packet() const {
    auto maybe_packet = WrapLegacyPacket<PayloadT>(
        holder_->context->InputSidePackets().Get(Tag(), Index()));
    // Framework guarantees correct type (PayloadT) of the side input.
    ABSL_CHECK_OK(maybe_packet);
    return std::move(maybe_packet).value();
  }
};

template <typename PayloadT>
class Output<ContextSpecializer, PayloadT>
    : public internal_port::Port<ContextSpecializer, OutputStreamField> {
 public:
  using Payload = PayloadT;
  using internal_port::Port<ContextSpecializer, OutputStreamField>::Port;

  void Send(const PayloadT& payload) const {
    holder_->context->Outputs()
        .Get(Tag(), Index())
        .AddPacket(mediapipe::MakePacket<PayloadT>(payload).At(
            holder_->context->InputTimestamp()));
  }

  void Send(PayloadT&& payload) const {
    holder_->context->Outputs()
        .Get(Tag(), Index())
        .AddPacket(
            mediapipe::MakePacket<PayloadT>(std::forward<PayloadT>(payload))
                .At(holder_->context->InputTimestamp()));
  }

  void Send(std::unique_ptr<PayloadT> payload) const {
    holder_->context->Outputs()
        .Get(Tag(), Index())
        .AddPacket(mediapipe::Adopt(payload.release())
                       .At(holder_->context->InputTimestamp()));
  }

  void Send(Packet<PayloadT> packet) const {
    holder_->context->Outputs()
        .Get(Tag(), Index())
        .AddPacket(std::move(packet).AsLegacyPacket());
  }

  Timestamp NextTimestampBound() const {
    return holder_->context->Outputs().Get(Tag(), Index()).NextTimestampBound();
  }

  void SetNextTimestampBound(Timestamp timestamp) {
    holder_->context->Outputs()
        .Get(Tag(), Index())
        .SetNextTimestampBound(timestamp);
  }

  bool IsClosed() const {
    return holder_->context->Outputs().Get(Tag(), Index()).IsClosed();
  }

  void Close() const {
    holder_->context->Outputs().Get(Tag(), Index()).Close();
  }
};

template <typename PayloadT>
class SideOutput<ContextSpecializer, PayloadT>
    : public internal_port::Port<ContextSpecializer, OutputSidePacketField> {
 public:
  using Payload = PayloadT;
  using internal_port::Port<ContextSpecializer, OutputSidePacketField>::Port;

  void Set(const PayloadT& payload) const {
    holder_->context->OutputSidePackets()
        .Get(Tag(), Index())
        .Set(mediapipe::MakePacket<PayloadT>(payload));
  }

  void Set(PayloadT&& payload) const {
    holder_->context->OutputSidePackets()
        .Get(Tag(), Index())
        .Set(mediapipe::MakePacket<PayloadT>(std::forward<PayloadT>(payload)));
  }

  void Set(Packet<PayloadT> packet) const {
    holder_->context->OutputSidePackets()
        .Get(Tag(), Index())
        .Set(std::move(packet).AsLegacyPacket());
  }
};

template <typename ProtoT>
class Options<ContextSpecializer, ProtoT> {
 public:
  using Field = OptionsField;
  using Specializer = ContextSpecializer;
  using Payload = ProtoT;

  const ProtoT& Get() const { return holder_->context->Options<ProtoT>(); }

  const ProtoT& operator()() const { return Get(); }

 protected:
  template <typename V, typename H>
  friend void internal_port::SetCalculatorContextHolder(V& v, H& holder);

  // Not owned, set by the framework.
  internal_port::CalculatorContextHolder* holder_ = nullptr;
};

///////////////////////// Implementation details below /////////////////////////

namespace internal {

template <class... F>
struct Overload : F... {
  using F::operator()...;
};
template <class... F>
explicit Overload(F...) -> Overload<F...>;

template <class T, class... U>
struct First {
  using type = T;
};

template <typename T, int&... DoNotSpecify, typename F>
auto VisitPacketOrDie(F&& visitor, const mediapipe::Packet& packet) {
  return std::forward<F>(visitor)(packet.Get<T>());
}

template <typename T, typename U, typename... Rest, int&... DoNotSpecify,
          typename F>
auto VisitPacketOrDie(F&& visitor, const mediapipe::Packet& packet) {
  if (packet.ValidateAsType<T>().ok()) {
    return std::forward<F>(visitor)(packet.Get<T>());
  } else {
    return VisitPacketOrDie<U, Rest...>(std::forward<F>(visitor), packet);
  }
}

template <typename T, int&... DoNotSpecify, typename F>
auto VisitPacketAsPacketOrDie(F&& visitor, const mediapipe::Packet& packet) {
  ABSL_CHECK_OK(packet.ValidateAsType<T>());
  return std::forward<F>(visitor)(WrapLegacyPacket<T>(packet).value());
}

template <typename T, typename U, typename... Rest, int&... DoNotSpecify,
          typename F>
auto VisitPacketAsPacketOrDie(F&& visitor, const mediapipe::Packet& packet) {
  if (packet.ValidateAsType<T>().ok()) {
    return std::forward<F>(visitor)(WrapLegacyPacket<T>(packet).value());
  } else {
    return VisitPacketAsPacketOrDie<U, Rest...>(std::forward<F>(visitor),
                                                packet);
  }
}

}  // namespace internal

template <typename... PayloadTs>
template <int&... DoNotSpecify, class... F,
          std::enable_if_t<(sizeof...(F) > 1), bool>>
auto Input<ContextSpecializer, OneOf<PayloadTs...>>::VisitOrDie(
    F&&... visitors) const {
  // Check that the number of functors matches the number of types.
  static_assert(sizeof...(F) == sizeof...(PayloadTs),
                "The number of provided visitors must match the number of "
                "types in the OneOf");
  // Check that the overload set `f` is callable with each type in PayloadTs.
  static_assert(
      (std::is_invocable_v<decltype(visitors), const PayloadTs&> && ...),
      "The provided visitors must be able to handle all types in the OneOf "
      "in the same order as types listed in the OneOf");
  return VisitOrDie(internal::Overload{std::forward<F>(visitors)...});
}

template <typename... PayloadTs>
template <int&... DoNotSpecify, class Visitor>
auto Input<ContextSpecializer, OneOf<PayloadTs...>>::VisitOrDie(
    Visitor&& visitor) const {
  // Check that the visitor is callable with each type in PayloadTs.
  static_assert(
      (std::is_invocable_v<Visitor, const PayloadTs&> && ...),
      "The provided visitor must be able to handle all types in the OneOf");

  using FirstT = typename internal::First<PayloadTs...>::type;
  using ResultType = std::invoke_result_t<Visitor, const FirstT&>;
  static_assert(
      (std::is_same_v<ResultType,
                      std::invoke_result_t<Visitor, const PayloadTs&>> &&
       ...),
      "All visitor overloads must have the same return type");
  return internal::VisitPacketOrDie<PayloadTs...>(
      std::forward<Visitor>(visitor),
      holder_->context->Inputs().Get(Tag(), Index()).Value());
}

template <typename... PayloadTs>
template <int&... DoNotSpecify, class Visitor>
auto Input<ContextSpecializer, OneOf<PayloadTs...>>::VisitAsPacketOrDie(
    Visitor&& visitor) const {
  // Check that the visitor is callable with each type in PayloadTs.
  static_assert(
      (std::is_invocable_v<Visitor, Packet<PayloadTs>> && ...),
      "The provided visitor must be able to handle all types in the OneOf");

  using FirstT = typename internal::First<PayloadTs...>::type;
  using ResultType = std::invoke_result_t<Visitor, Packet<FirstT>>;
  static_assert(
      (std::is_same_v<ResultType,
                      std::invoke_result_t<Visitor, Packet<PayloadTs>>> &&
       ...),
      "All visitor overloads must have the same return type");
  return internal::VisitPacketAsPacketOrDie<PayloadTs...>(
      std::forward<Visitor>(visitor),
      holder_->context->Inputs().Get(Tag(), Index()).Value());
}

}  // namespace mediapipe::api3

#endif  // MEDIAPIPE_FRAMEWORK_API3_CALCULATOR_CONTEXT_H_
